﻿from simcell.genetic.amino import AminoFactory
from simcell.genetic.protein import Protein

class GeneticCode:
    """Genetic Code"""

    TRANSLATION_TABLE = {'aa': AminoFactory.H,
                         'au': AminoFactory.N,
                         'uu': AminoFactory.V}

    START_CODON = 'aa'
    STOP_CODON  = 'ua'
    LEN_CODON = 2
          
    @classmethod
    def isStartCodon(cls, codon):
        return codon == cls.START_CODON
    
    @classmethod  
    def isStopCodon(cls, codon):
        return codon == cls.STOP_CODON

    @ classmethod
    def translateRna(cls, rna):
        rna_strand = "".join(map(str, rna.strand))        
        protein = Protein()
        translating = False
        for i in xrange(0, len(rna_strand), GeneticCode.LEN_CODON):
            codon = rna_strand[i:i+GeneticCode.LEN_CODON]
            if len(codon) == GeneticCode.LEN_CODON:
                if translating:
                    if cls.isStopCodon(codon):
                        return protein
                    else:
                        amino = cls.TRANSLATION_TABLE[codon]
                        protein.addAmino(amino)
                        
                else:
                    if cls.isStartCodon(codon):
                        amino = cls.TRANSLATION_TABLE[codon]
                        protein.addAmino(amino)
                        translating = True
                        
        return protein

    @classmethod
    def codonsToProtein(cls, seq):
        translating = False
        proteina = []
        for codon in seq:
            if translating:
                if self.isStopCodon(codon):
                    return proteina
                else:
                    amino = self.translateCodon(codon)
                    if not amino is None:
                      proteina.append(amino)
            else:
                if self.isStartCodon(codon):
                    proteina.append(self.translateCodon(codon))
                    translating = True
        
        if len(proteina) == 0:
          proteina = None
          
        return proteina
